"""
Extension Router for GSM Extension System

This module provides routing functionality to direct commands and configurations
to appropriate extensions based on various criteria.
"""

from typing import Dict, List, Optional, Any, Callable, Tuple
from pathlib import Path
import logging
import re
import xml.etree.ElementTree as ET

from .extension_registry import ExtensionRegistry


class RoutingRule:
    """A rule for routing commands/configs to extensions."""
    
    def __init__(self, name: str, priority: int = 50):
        """
        Initialize routing rule.
        
        Args:
            name: Rule name
            priority: Priority (lower numbers have higher priority)
        """
        self.name = name
        self.priority = priority
    
    def matches(self, context: Dict[str, Any]) -> bool:
        """
        Check if rule matches the given context.
        
        Args:
            context: Routing context containing command, config, etc.
            
        Returns:
            bool: True if rule matches
        """
        raise NotImplementedError
    
    def get_extension_name(self, context: Dict[str, Any]) -> Optional[str]:
        """
        Get extension name for matched context.
        
        Args:
            context: Routing context
            
        Returns:
            Optional[str]: Extension name or None
        """
        raise NotImplementedError


class FilePatternRule(RoutingRule):
    """Routes based on file path patterns."""
    
    def __init__(self, name: str, patterns: List[str], extension_name: str, 
                 priority: int = 50):
        """
        Initialize file pattern rule.
        
        Args:
            name: Rule name
            patterns: List of file path patterns (supports glob patterns)
            extension_name: Target extension name
            priority: Rule priority
        """
        super().__init__(name, priority)
        self.patterns = [re.compile(p.replace('*', '.*').replace('?', '.')) for p in patterns]
        self.extension_name = extension_name
    
    def matches(self, context: Dict[str, Any]) -> bool:
        """Check if config path matches patterns."""
        config_path = context.get('config_path')
        if not config_path:
            return False
        
        path_str = str(config_path)
        return any(pattern.search(path_str) for pattern in self.patterns)
    
    def get_extension_name(self, context: Dict[str, Any]) -> Optional[str]:
        """Return target extension name."""
        return self.extension_name


class XMLContentRule(RoutingRule):
    """Routes based on XML configuration content."""
    
    def __init__(self, name: str, xpath_conditions: Dict[str, str], 
                 extension_name: str, priority: int = 50):
        """
        Initialize XML content rule.
        
        Args:
            name: Rule name
            xpath_conditions: Dictionary of xpath -> expected value
            extension_name: Target extension name
            priority: Rule priority
        """
        super().__init__(name, priority)
        self.xpath_conditions = xpath_conditions
        self.extension_name = extension_name
    
    def matches(self, context: Dict[str, Any]) -> bool:
        """Check if XML content matches conditions."""
        config_path = context.get('config_path')
        if not config_path or not config_path.exists():
            return False
        
        try:
            # Parse XML
            tree = ET.parse(config_path)
            root = tree.getroot()
            
            # Check all conditions
            for xpath, expected_value in self.xpath_conditions.items():
                elements = root.findall(xpath)
                if not elements:
                    return False
                
                # Check if any element matches expected value
                matches = False
                for elem in elements:
                    if elem.text and expected_value.lower() in elem.text.lower():
                        matches = True
                        break
                    # Also check attributes
                    for attr_value in elem.attrib.values():
                        if expected_value.lower() in str(attr_value).lower():
                            matches = True
                            break
                    if matches:
                        break
                
                if not matches:
                    return False
            
            return True
            
        except Exception:
            return False
    
    def get_extension_name(self, context: Dict[str, Any]) -> Optional[str]:
        """Return target extension name."""
        return self.extension_name


class CommandRule(RoutingRule):
    """Routes based on command name and arguments."""
    
    def __init__(self, name: str, commands: List[str], extension_name: str, 
                 priority: int = 50):
        """
        Initialize command rule.
        
        Args:
            name: Rule name  
            commands: List of command names to match
            extension_name: Target extension name
            priority: Rule priority
        """
        super().__init__(name, priority)
        self.commands = commands
        self.extension_name = extension_name
    
    def matches(self, context: Dict[str, Any]) -> bool:
        """Check if command matches."""
        command = context.get('command')
        return command in self.commands if command else False
    
    def get_extension_name(self, context: Dict[str, Any]) -> Optional[str]:
        """Return target extension name."""
        return self.extension_name


class ExtensionRouter:
    """
    Routes commands and configurations to appropriate extensions.
    
    The router uses a set of configurable rules to determine which extension
    should handle a particular command or configuration file.
    """
    
    def __init__(self, registry: Optional[ExtensionRegistry] = None, 
                 logger: Optional[logging.Logger] = None):
        """
        Initialize extension router.
        
        Args:
            registry: Extension registry instance
            logger: Optional logger instance
        """
        self._registry = registry or ExtensionRegistry()
        self._logger = logger or logging.getLogger(__name__)
        
        # Routing rules (sorted by priority)
        self._rules: List[RoutingRule] = []
        
        # Default extension for fallback
        self._default_extension: Optional[str] = None
        
        # Caching for performance
        self._routing_cache: Dict[str, str] = {}
        self._cache_max_size = 100
        
        # Initialize default rules
        self._setup_default_rules()
        
        self._logger.debug("Extension router initialized")
    
    def add_rule(self, rule: RoutingRule) -> None:
        """
        Add a routing rule.
        
        Args:
            rule: Routing rule to add
        """
        self._rules.append(rule)
        self._rules.sort(key=lambda r: r.priority)
        
        # Clear cache when rules change
        self._routing_cache.clear()
        
        self._logger.debug(f"Added routing rule: {rule.name} (priority: {rule.priority})")
    
    def remove_rule(self, rule_name: str) -> bool:
        """
        Remove a routing rule by name.
        
        Args:
            rule_name: Name of rule to remove
            
        Returns:
            bool: True if rule was removed
        """
        for i, rule in enumerate(self._rules):
            if rule.name == rule_name:
                del self._rules[i]
                self._routing_cache.clear()
                self._logger.debug(f"Removed routing rule: {rule_name}")
                return True
        return False
    
    def set_default_extension(self, extension_name: str) -> None:
        """
        Set default extension for fallback routing.
        
        Args:
            extension_name: Name of default extension
        """
        self._default_extension = extension_name
        self._logger.debug(f"Set default extension: {extension_name}")
    
    def route_command(self, command: str, args: Optional[Dict[str, Any]] = None) -> Optional[str]:
        """
        Route command to appropriate extension.
        
        Args:
            command: Command name
            args: Optional command arguments
            
        Returns:
            Optional[str]: Extension name that should handle the command
        """
        context = {
            'command': command,
            'args': args or {},
        }
        
        # Check cache first
        cache_key = f"cmd:{command}"
        if cache_key in self._routing_cache:
            extension_name = self._routing_cache[cache_key]
            self._logger.debug(f"Using cached routing for command {command}: {extension_name}")
            return extension_name
        
        # Apply routing rules
        extension_name = self._apply_rules(context)
        
        # Cache result
        self._cache_result(cache_key, extension_name)
        
        self._logger.debug(f"Routed command {command} to extension: {extension_name}")
        return extension_name
    
    def route_config(self, config_path: Path) -> Optional[str]:
        """
        Route configuration file to appropriate extension.
        
        Args:
            config_path: Path to configuration file
            
        Returns:
            Optional[str]: Extension name that should handle the config
        """
        context = {
            'config_path': config_path,
        }
        
        # Check cache first
        cache_key = f"cfg:{str(config_path)}"
        if cache_key in self._routing_cache:
            extension_name = self._routing_cache[cache_key]
            self._logger.debug(f"Using cached routing for config {config_path}: {extension_name}")
            return extension_name
        
        # Apply routing rules
        extension_name = self._apply_rules(context)
        
        # If no rule matched, try extension can_handle_config methods
        if not extension_name:
            extension_name = self._try_extension_handlers(config_path)
        
        # Cache result
        self._cache_result(cache_key, extension_name)
        
        self._logger.debug(f"Routed config {config_path} to extension: {extension_name}")
        return extension_name
    
    def get_routing_info(self) -> Dict[str, Any]:
        """
        Get routing configuration information.
        
        Returns:
            Dict[str, Any]: Routing information
        """
        return {
            'total_rules': len(self._rules),
            'rules': [
                {
                    'name': rule.name,
                    'type': type(rule).__name__,
                    'priority': rule.priority,
                }
                for rule in self._rules
            ],
            'default_extension': self._default_extension,
            'cache_size': len(self._routing_cache),
            'cache_max_size': self._cache_max_size,
        }
    
    def clear_cache(self) -> None:
        """Clear routing cache."""
        self._routing_cache.clear()
        self._logger.debug("Cleared routing cache")
    
    def _setup_default_rules(self) -> None:
        """Setup default routing rules."""
        
        # Hybrid PETSc-Kokkos extension rules
        hybrid_patterns = [
            '*/prototypes/hybrid/*',
            '*/hybrid/*',
            '*hybrid*.xml',
        ]
        self.add_rule(FilePatternRule(
            name="hybrid-path-patterns",
            patterns=hybrid_patterns,
            extension_name="hybrid-petsc-kokkos",
            priority=10
        ))
        
        # XML content-based routing for hybrid
        self.add_rule(XMLContentRule(
            name="hybrid-xml-content", 
            xpath_conditions={
                ".//solver_type": "hybrid",
                ".//framework": "hybrid", 
                ".//backend": "petsc-kokkos",
            },
            extension_name="hybrid-petsc-kokkos",
            priority=20
        ))
        
        # Original GCR-NCCL patterns
        original_patterns = [
            '*gcr*.xml',
            '*nccl*.xml',
            '*/original/*',
        ]
        self.add_rule(FilePatternRule(
            name="original-path-patterns",
            patterns=original_patterns,
            extension_name="original-gcr-nccl",
            priority=30
        ))
        
        self._logger.debug("Setup default routing rules")
    
    def _apply_rules(self, context: Dict[str, Any]) -> Optional[str]:
        """Apply routing rules to context."""
        for rule in self._rules:
            try:
                if rule.matches(context):
                    extension_name = rule.get_extension_name(context)
                    if extension_name and self._registry.get_extension(extension_name):
                        return extension_name
            except Exception as e:
                self._logger.warning(f"Error applying rule {rule.name}: {str(e)}")
        
        return self._default_extension
    
    def _try_extension_handlers(self, config_path: Path) -> Optional[str]:
        """Try extension can_handle_config methods."""
        # Get extensions in load order for consistent behavior
        for extension_name in self._registry.get_load_order():
            extension = self._registry.get_extension(extension_name)
            if extension:
                try:
                    if extension.can_handle_config(config_path):
                        return extension_name
                except Exception as e:
                    self._logger.warning(f"Error checking config compatibility for {extension_name}: {str(e)}")
        
        return None
    
    def _cache_result(self, cache_key: str, extension_name: Optional[str]) -> None:
        """Cache routing result."""
        if extension_name:
            # Manage cache size
            if len(self._routing_cache) >= self._cache_max_size:
                # Remove oldest entry (simple FIFO)
                oldest_key = next(iter(self._routing_cache))
                del self._routing_cache[oldest_key]
            
            self._routing_cache[cache_key] = extension_name